from langchain_core.output_parsers import StrOutputParser
from typing import TypedDict
from langchain_openai import ChatOpenAI
from config import COMMON_LLM_MODEL_NAME, COMMON_LLM_URL
from langgraph.graph import END, StateGraph
from langchain_community.chat_models import ChatZhipuAI
# from plan_class import *
from survey_prompt_a import *
from survey_prompt_b import *
from survey_prompt_c import *
from survey_prompt_d import *
from survey_prompt_a_EN import *
from survey_prompt_b_EN import *
from survey_prompt_c_EN import *
from survey_prompt_d_EN import *
import re
from system_prompt import *
from system_prompt_EN import *
import pickle
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
import argparse
from collections import defaultdict
import os
import itertools
import time
from zhipuai import ZhipuAI
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--topic', type=str, required=True,
                       help='Research topic for paper review')
    return parser.parse_args()

### 决策树 ###
#                                            用户输入
# a.对一个技术概念的调研综述  b.对一个方向研究现状的综述  c.对多篇文献的对比分析综述  d.对一个技术方法研究脉络的综述
# 1.引言                  1.引言                   1.引言                  1.引言
# 2.基础概念和定义          2.基础概念和定义           2.对比历史工作           2.基础概念和定义
# 3.技术详解（历史和相关工作）3.研究现状                3.对比现状               3.历史工作
# 4.实际应用               4.应用现状                4.对比实际应用           4.研究现状
# 5.挑战和未来工作          5.总结                   5.总结                  5.挑战和未来工作
# 6.总结                                                                 6.总结

class Agent(TypedDict):
    user_input: str
    language: int
    theme: str
    introduction: str
    history: str
    tech: str
    app: str
    future: str
    sum: str
    definition_paper: str
    importance_paper: str
    background_paper: str
    present_work_paper: str
    related_work_paper: str
    application_paper: str
    challenge_paper: str
    future_work_paper: str
    survey: str
    step1_a: str
    step2_a: str
    step3_a: str
    step4_a: str
    step5_a: str
    step6_a: str
    step1_b: str
    step2_b: str
    step3_b: str
    step4_b: str
    step5_b: str
    step1_c: str
    step2_c: str
    step3_c: str
    step4_c: str
    step5_c: str
    step1_d: str
    step2_d: str
    step3_d: str
    step4_d: str
    step5_d: str
    step6_d: str
    paper_count: int
    paper_list: list


# 配置大模型
#llm = ChatOpenAI(base_url=COMMON_LLM_URL, model=COMMON_LLM_MODEL_NAME, api_key='anything')

llm = ChatOpenAI(base_url='https://open.bigmodel.cn/api/paas/v4', model='glm-4-plus', api_key='6de473a119cc44718d305f7fcf5b1dba.SmdMLTsAnA3QJI2h', temperature=0)
# 设置初始语言，默认0为中文，1为英文
# 通过第一个字符判断用户的提问是中文还是英文
# def judge_language(state):
#     user_input = state['user_input']
#     if user_input[0].isalpha():
#         return {'language': 1}
#     else:
#         return {'language': 0}


def get_paper(query, top_k=20):
    url = "http://180.184.65.98:38880/atomgit/search_papers"

    # 设置请求参数
    params = {
        'query': query,
        'top_k': top_k
    }
    paper = []
    try:
        # 发起 GET 请求
        response = requests.get(url, params=params)

        # 如果响应状态码是 200，表示请求成功
        if response.status_code == 200:
            # 将响应内容解析为 JSON
            papers = response.json()
            for p in papers:
                p['entity']['chunk_text'] = p['entity']['chunk_text'][:2500]
                paper.append(p)
            return paper
        else:
            print(f"Error: Received status code {response.status_code}")
            return None
    except requests.RequestException as e:
        # 捕获请求异常
        print(f"An error occurred: {e}")
        return None


# 获得paper之后，使用大模型总结，压缩文本
# 并发预处理，加快速度
# def process_paper(papers, user_input, sum_paper_llm, section):
#     result = []
#     for paper in papers:
#         #print(paper)
#         paper['entity']['chunk_text'] = sum_paper_llm.invoke(
#                 {'user_input': user_input, 'paper_content': paper['entity']['chunk_text']}
#             )
#         result.append(paper)
#     return {
#         'section': section,
#         'result': result
#     }
def process_single_paper(paper, user_input, sum_paper_llm):
    return {
        'paper_title': paper['entity']['paper_title'],
        'chunk_id': paper['entity']['chunk_id'],
        'chunk_text': sum_paper_llm.invoke({
            'user_input': user_input,
            'paper_content': paper['entity']['chunk_text']
        })
    }


def process_paper(all_papers, user_input, sum_paper_llm):
    from concurrent.futures import ThreadPoolExecutor
    # 创建扁平化任务列表
    tasks = []
    for section in all_papers:
        tasks.extend([
            (paper, section)
            for paper in all_papers[section]
        ])

    # 三级并发：论文级别
    with ThreadPoolExecutor(max_workers=os.cpu_count() * 2) as executor:
        futures = [
            executor.submit(
                process_single_paper,
                paper,
                user_input,
                sum_paper_llm
            ) for (paper, section) in tasks
        ]

        # 动态分配线程资源
        results = defaultdict(list)
        for future, (_, section) in zip(as_completed(futures), tasks):
            results[section].append(future.result())

    return results

# def get_all_papers(state):
#     theme_llm = theme_template | llm | StrOutputParser()
#     user_input = state['user_input']
#     theme = theme_llm.invoke({'user_input': user_input})
#     # 把文章切分成多个章节，每个章节根据文章主题来获得一定数量的引用论文片段
#     sections = ['definition', 'importance', 'background', 'present_work', 'related_work', 'application', 'challenge',
#                 'future_work']
#     # Queries for different sections
#     queries = {
#         'definition': ["definition of " + theme],
#         'importance': ["importance of " + theme],
#         'background': ["background of " + theme],
#         'present_work': ["present_work of " + theme],
#         'related_work': ["related_work of " + theme],
#         'application': ["application of " + theme],
#         'challenge': ["challenge of " + theme],
#         'future_work': ["future_work of " + theme]
#     }
#     # Fetch papers for each query
#     papers = {section: get_paper(queries[section], top_k=15) for section in sections}
#     results = {
#         'definition_paper': [],
#         'importance_paper': [],
#         'background_paper': [],
#         'present_work_paper': [],
#         'related_work_paper': [],
#         'application_paper': [],
#         'challenge_paper': [],
#         'future_work_paper': []
#     }
#     sum_paper_llm = sum_paper_template_EN | llm | StrOutputParser()
#     for char in user_input:
#         if '\u4e00' <= char <= '\u9fa5':
#             sum_paper_llm = sum_paper_template | llm | StrOutputParser()
#
#
#     with ThreadPoolExecutor() as executor:
#         futures = []
#         for section in sections:
#             futures.append(executor.submit(process_paper, papers[section], user_input, sum_paper_llm, section))
#
#         for future in as_completed(futures):
#             result = future.result()
#             section = result['section']
#             results[section + '_paper'].append(result['result'])
#     # print(len(results['definition_paper']))
#     # print(len(results['position_paper']))
#     # print(len(results['history_paper']))
#     # print(len(results['position_paper']))
#     # print(len(results['tech_paper']))
#     # print(len(results['app_paper']))
#     # print(len(results['future_paper']))
#     # print(len(results['sum_paper']))
#     definition_paper = [{'paper_title': i['entity']['paper_title'], 'chunk_id': i['entity']['chunk_id'], 'chunk_text': i['entity']['chunk_text']}for i in
#                         results['definition_paper'][0]]
#     importance_paper = [{'paper_title': i['entity']['paper_title'], 'chunk_id': i['entity']['chunk_id'], 'chunk_text': i['entity']['chunk_text']}for i in
#                         results['importance_paper'][0]]
#     background_paper = [{'paper_title': i['entity']['paper_title'], 'chunk_id': i['entity']['chunk_id'], 'chunk_text': i['entity']['chunk_text']}for i in
#                         results['background_paper'][0]]
#     present_work_paper = [{'paper_title': i['entity']['paper_title'], 'chunk_id': i['entity']['chunk_id'], 'chunk_text': i['entity']['chunk_text']}for i in
#                         results['present_work_paper'][0]]
#     related_work_paper = [{'paper_title': i['entity']['paper_title'], 'chunk_id': i['entity']['chunk_id'], 'chunk_text': i['entity']['chunk_text']}for i in
#                         results['related_work_paper'][0]]
#     application_paper = [{'paper_title': i['entity']['paper_title'], 'chunk_id': i['entity']['chunk_id'], 'chunk_text': i['entity']['chunk_text']}for i in
#                         results['application_paper'][0]]
#     challenge_paper = [{'paper_title': i['entity']['paper_title'], 'chunk_id': i['entity']['chunk_id'], 'chunk_text': i['entity']['chunk_text']}for i in
#                         results['challenge_paper'][0]]
#     future_work_paper = [{'paper_title': i['entity']['paper_title'], 'chunk_id': i['entity']['chunk_id'], 'chunk_text': i['entity']['chunk_text']}for i in
#                         results['future_work_paper'][0]]
#     return {'definition_paper': definition_paper, 'importance_paper': importance_paper,
#             'background_paper': background_paper, 'present_work_paper': present_work_paper,
#             'related_work_paper': related_work_paper, 'application_paper': application_paper,
#             'challenge_paper': challenge_paper, 'future_work_paper': future_work_paper}
def get_all_papers(state):
    # 1. 定义主题模板和初始化变量
    theme_llm = theme_template | llm | StrOutputParser()
    user_input = state['user_input']
    theme = theme_llm.invoke({'user_input': user_input})

    # 2. 定义章节和查询
    sections = ['definition', 'importance', 'background', 'present_work',
                'related_work', 'application', 'challenge', 'future_work']
    queries = {s: [f"{s} of {theme}"] for s in sections}

    # 3. 优化中文检测
    has_chinese = any('\u4e00' <= char <= '\u9fff' for char in user_input)
    sum_paper_llm = (sum_paper_template if has_chinese else sum_paper_template_EN) | llm | StrOutputParser()

    # 4. 并发获取所有章节的论文（第一级并发）
    papers_dict = {}
    with ThreadPoolExecutor(max_workers=os.cpu_count() * 2) as executor:
        # 创建所有章节的请求任务
        future_to_section = {
            executor.submit(get_paper, queries[section], 15): section
            for section in sections
        }
        # 收集结果
        for future in as_completed(future_to_section):
            section = future_to_section[future]
            papers_dict[section] = future.result() or []

    # 5. 创建所有论文的扁平化列表并处理（第二级并发）
    processed = defaultdict(list)
    with ThreadPoolExecutor(max_workers=os.cpu_count() * 2) as executor:  # 使用双倍核心数的线程
        # 准备所有论文处理任务
        futures = []
        for section in sections:
            for paper in papers_dict.get(section, []):
                futures.append(
                    executor.submit(
                        lambda p: {
                            'paper_title': p['entity']['paper_title'],
                            'chunk_id': p['entity']['chunk_id'],
                            'year': p['entity']['year'],
                            'chunk_text': sum_paper_llm.invoke({
                                'user_input': user_input,
                                'paper_content': p['entity']['chunk_text']
                            })
                        },
                        paper
                    )
                )

        # 按章节收集处理结果
        for future, section in zip(futures, itertools.chain(*[[s] * len(papers_dict[s]) for s in sections])):
            processed[section].append(future.result())

    # 6. 构建最终结果
    return {
        'definition_paper': processed.get('definition', []),
        'importance_paper': processed.get('importance', []),
        'background_paper': processed.get('background', []),
        'present_work_paper': processed.get('present_work', []),
        'related_work_paper': processed.get('related_work', []),
        'application_paper': processed.get('application', []),
        'challenge_paper': processed.get('challenge', []),
        'future_work_paper': processed.get('future_work', [])
    }

# 判断综述需求的类型
def judge_survey_type(state):
    judge_survey_llm = judge_survey_template_EN | llm | StrOutputParser()
    result = judge_survey_llm.invoke(
        {'user_input': state['user_input']}
    )
    print(result,'------')
    if 'a' in result.splitlines()[-1]:
        return 'a'
    elif 'b' in result.splitlines()[-1]:
        return 'b'
    elif 'c' in result.splitlines()[-1]:
        return 'c'
    else:
        return 'd'


# 确定综述关键词
def get_survey_theme(state):
    user_input = state['user_input']
    if user_input[0].isalpha():
        theme_llm = theme_template_EN | llm | StrOutputParser()
    else:
        theme_llm = theme_template | llm | StrOutputParser()
    theme = theme_llm.invoke(
        {'user_input': user_input}
    )
    return {'theme': theme}

# 该函数处理文章最后的引用，使其符合提交标准
def process_article(article):
    # 创建一个字典，便于快速查找检索id和对应的序号
    sup_pattern = re.compile(r'<sup>(.*?)</sup>')

    # 初始化替换后的文章和记录信息的列表
    replaced_article = article
    paper_info_list = []

    # 使用一个计数器来生成序号
    count = 1

    # 查找所有的 <sup>...</sup> 并进行替换
    for match in re.finditer(sup_pattern, article):
        # 提取 <sup> 标签中的内容
        original_content = match.group(1)
        # 提取 paper_title 和 chunk_id
        if '***' in original_content:
            try:
                paper_title, chunk_id, year = original_content.split('***')
                paper_title = paper_title.strip()  # 去除前后的空格
                chunk_id = chunk_id.strip()
                year = year.strip() # 去除前后的空格
            except Exception as e:
                paper_title = original_content
                chunk_id = ''
                year = ''
        else:
            paper_title = original_content
            chunk_id = ''
            year = ''

        # 替换为序号
        replaced_article = replaced_article.replace(match.group(0), f'<sup>{count}</sup>', 1)

        # 将提取的信息保存到列表
        paper_info_list.append(f'[{count}] {paper_title}, {year}, chunk {chunk_id}')

        # 增加序号
        count += 1

    # 在文章末尾添加所有的参考信息
    replaced_article += '\n' + '\n'.join(paper_info_list)

    return replaced_article





# a类综述的第一步
def write_step1_a(state):
    # print(len(state['definition_paper']))
    # print(len(state['importance_paper']))
    # print(len(state['background_paper']))
    # print(len(state['present_work_paper']))
    # print(len(state['related_work_paper']))
    # print(len(state['application_paper']))
    # print(len(state['challenge_paper']))
    # print(len(state['future_work_paper']))
    user_input = state['user_input']
    background_paper = state['background_paper']
    importance_paper = state['importance_paper']
    write_step1_llm = write_step1_template_a_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step1_llm = write_step1_template_a | llm | StrOutputParser()
    step1_a = write_step1_llm.invoke(
        {'user_input': user_input,
         'background_paper': background_paper,
         'importance_paper': importance_paper}
    )
    # 记录每篇论文的引用次数，最多为2次
    return {'step1_a': step1_a}


# a类综述的第二步
def write_step2_a(state):
    user_input = state['user_input']
    step1_a = state['step1_a']
    definition_paper = state['definition_paper']
    write_step2_llm = write_step2_template_a_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step2_llm = write_step2_template_a | llm | StrOutputParser()
    step2_a = write_step2_llm.invoke(
        {'user_input': user_input,
         'step1_a': step1_a,
         'definition_paper': definition_paper}
    )
    return {'step2_a': step2_a}


# a类综述的第三步
def write_step3_a(state):
    user_input = state['user_input']
    present_work_paper = state['present_work_paper']
    related_work_paper = state['related_work_paper']
    write_step3_llm = write_step3_template_a_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step3_llm = write_step3_template_a | llm | StrOutputParser()
    step3_a = write_step3_llm.invoke(
        {'user_input': user_input,
         'present_work_paper': present_work_paper,
         'related_work_paper': related_work_paper}
    )
    return {'step3_a': step3_a}


# a类综述的第四步
def write_step4_a(state):
    user_input = state['user_input']
    application_paper = state['application_paper']
    write_step4_llm = write_step4_template_a_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step4_llm = write_step4_template_a | llm | StrOutputParser()
    step4_a = write_step4_llm.invoke(
        {'user_input': user_input,
         'application_paper': application_paper}
    )
    return {'step4_a': step4_a}


# a类综述的第五步
def write_step5_a(state):
    user_input = state['user_input']
    challenge_paper = state['challenge_paper']
    future_work_paper = state['future_work_paper']
    write_step5_llm = write_step5_template_a_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step5_llm = write_step5_template_a | llm | StrOutputParser()
    step5_a = write_step5_llm.invoke(
        {'user_input': user_input,
         'challenge_paper': challenge_paper,
         'future_work_paper': future_work_paper}
    )
    return {'step5_a': step5_a}


# a类综述的第六步
def write_step6_a(state):
    if state['step1_a'] and state['step2_a'] and state['step3_a'] and state['step4_a'] and state['step5_a']:
        user_input = state['user_input']
        finished_work = state['step1_a']+'\n'+state['step2_a']+'\n'+state['step3_a']+'\n'+state['step4_a']+'\n'+state['step5_a']+'\n'
        write_step6_llm = write_step6_template_a_EN | llm | StrOutputParser()
        for char in user_input:
            if '\u4e00' <= char <= '\u9fa5':
                write_step6_llm = write_step6_template_a | llm | StrOutputParser()
        step6_a = write_step6_llm.invoke(
            {'user_input': user_input,
             'finished_work': finished_work}
        )
        article = process_article(finished_work+step6_a)
        with open('review.md', 'w', encoding='utf-8') as file:
            # 将传入的字符串内容写入文件
            file.write(article)
        return {'step6_a': step6_a}


# b类综述的第一步
def write_step1_b(state):
    user_input = state['user_input']
    background_paper = state['background_paper']
    importance_paper = state['importance_paper']
    write_step1_llm = write_step1_template_b_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step1_llm = write_step1_template_b | llm | StrOutputParser()
    step1_b = write_step1_llm.invoke(
        {'user_input': user_input,
         'background_paper': background_paper,
         'importance_paper': importance_paper}
    )
    return {'step1_b': step1_b}


# b类综述的第二步
def write_step2_b(state):
    user_input = state['user_input']
    step1_b = state['step1_b']
    definition_paper = state['definition_paper']
    write_step2_llm = write_step2_template_b_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step2_llm = write_step2_template_b | llm | StrOutputParser()
    step2_b = write_step2_llm.invoke(
        {'user_input': user_input,
         'step1_b': step1_b,
         'definition_paper': definition_paper}
    )
    return {'step2_b': step2_b}


# b类综述的第三步
def write_step3_b(state):
    user_input = state['user_input']
    related_work_paper = state['related_work_paper']
    write_step3_llm = write_step3_template_b_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step3_llm = write_step3_template_b | llm | StrOutputParser()
    step3_b = write_step3_llm.invoke(
        {'user_input': user_input,
         'related_work_paper': related_work_paper}
    )
    return {'step3_b': step3_b}


# b类综述的第四步
def write_step4_b(state):
    user_input = state['user_input']
    application_paper = state['application_paper']
    write_step4_llm = write_step4_template_b_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step4_llm = write_step4_template_b | llm | StrOutputParser()
    step4_b = write_step4_llm.invoke(
        {'user_input': user_input,
         'application_paper': application_paper}
    )
    return {'step4_b': step4_b}


# b类综述的第五步
def write_step5_b(state):
    user_input = state['user_input']
    finished_work = state['step1_b']+'\n'+state['step2_b']+'\n'+state['step3_b']+'\n'+state['step4_b']+'\n'
    write_step5_llm = write_step5_template_b_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step5_llm = write_step5_template_b | llm | StrOutputParser()
    step5_b = write_step5_llm.invoke(
        {'user_input': user_input,
         'finished_work': finished_work}
    )
    article = process_article(finished_work + step5_b)
    with open('review.md', 'w', encoding='utf-8') as file:
        # 将传入的字符串内容写入文件
        file.write(article)
    return {'step5_b': step5_b}


# c类综述的第一步
def write_step1_c(state):
    user_input = state['user_input']
    background_paper = state['background_paper']
    importance_paper = state['importance_paper']
    write_step1_llm = write_step1_template_c_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step1_llm = write_step1_template_c | llm | StrOutputParser()
    step1_c = write_step1_llm.invoke(
        {'user_input': user_input,
         'background_paper': background_paper,
         'importance_paper': importance_paper}
    )
    return {'step1_c': step1_c}


# c类综述的第二步
def write_step2_c(state):
    user_input = state['user_input']
    step1_c = state['step1_c']
    present_work_paper = state['present_work_paper']
    write_step2_llm = write_step2_template_c_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step2_llm = write_step2_template_c | llm | StrOutputParser()
    step2_c = write_step2_llm.invoke(
        {'user_input': user_input,
         'step1_c': step1_c,
         'present_work_paper': present_work_paper}
    )
    return {'step2_c': step2_c}


# c类综述的第三步
def write_step3_c(state):
    user_input = state['user_input']
    related_work_paper = state['related_work_paper']
    write_step3_llm = write_step3_template_c_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step3_llm = write_step3_template_c | llm | StrOutputParser()
    step3_c = write_step3_llm.invoke(
        {'user_input': user_input,
         'related_work_paper': related_work_paper}
    )
    return {'step3_c': step3_c}


# c类综述的第四步
def write_step4_c(state):
    user_input = state['user_input']
    application_paper = state['application_paper']
    write_step4_llm = write_step4_template_c_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step4_llm = write_step4_template_c | llm | StrOutputParser()
    step4_c = write_step4_llm.invoke(
        {'user_input': user_input,
         'application_paper': application_paper}
    )
    return {'step4_c': step4_c}


# c类综述的第五步
def write_step5_c(state):
    user_input = state['user_input']
    finished_work = state['step1_c']+'\n'+state['step2_c']+'\n'+state['step3_c']+'\n'+state['step4_c']+'\n'
    write_step5_llm = write_step5_template_c_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step5_llm = write_step5_template_c | llm | StrOutputParser()
    step5_c = write_step5_llm.invoke(
        {'user_input': user_input,
         'finished_work': finished_work}
    )
    article = process_article(finished_work + step5_c)
    with open('review.md', 'w', encoding='utf-8') as file:
        # 将传入的字符串内容写入文件
        file.write(article)
    return {'step5_c': step5_c}


# d类综述的第一步
def write_step1_d(state):
    user_input = state['user_input']
    background_paper = state['background_paper']
    importance_paper = state['importance_paper']
    write_step1_llm = write_step1_template_d_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step1_llm = write_step1_template_d | llm | StrOutputParser()
    step1_d = write_step1_llm.invoke(
        {'user_input': user_input,
         'background_paper': background_paper,
         'importance_paper': importance_paper}
    )
    return {'step1_d': step1_d}


# d类综述的第二步
def write_step2_d(state):
    user_input = state['user_input']
    step1_d = state['step1_d']
    definition_paper = state['definition_paper']
    write_step2_llm = write_step2_template_d_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step2_llm = write_step2_template_d | llm | StrOutputParser()
    step2_d = write_step2_llm.invoke(
        {'user_input': user_input,
         'step1_d': step1_d,
         'definition_paper': definition_paper}
    )
    return {'step2_d': step2_d}


# d类综述的第三步
def write_step3_d(state):
    user_input = state['user_input']
    present_work_paper = state['present_work_paper']
    write_step3_llm = write_step3_template_d_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step3_llm = write_step3_template_d | llm | StrOutputParser()
    step3_d = write_step3_llm.invoke(
        {'user_input': user_input,
         'present_work_paper': present_work_paper}
    )
    return {'step3_d': step3_d}


# d类综述的第四步
def write_step4_d(state):
    user_input = state['user_input']
    related_work_paper = state['related_work_paper']
    write_step4_llm = write_step4_template_d_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step4_llm = write_step4_template_d | llm | StrOutputParser()
    step4_d = write_step4_llm.invoke(
        {'user_input': user_input,
         'related_work_paper': related_work_paper}
    )
    return {'step4_d': step4_d}


# d类综述的第五步
def write_step5_d(state):
    user_input = state['user_input']
    challenge_paper = state['challenge_paper']
    future_work_paper = state['future_work_paper']
    write_step5_llm = write_step5_template_d_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step5_llm = write_step5_template_d | llm | StrOutputParser()
    step5_d = write_step5_llm.invoke(
        {'challenge_paper': challenge_paper,
         'future_work_paper': future_work_paper,
         'user_input': user_input}
    )
    return {'step5_d': step5_d}


# d类综述的第六步
def write_step6_d(state):
    user_input = state['user_input']
    finished_work = state['step1_d']+'\n'+state['step2_d']+'\n'+state['step3_d']+'\n'+state['step4_d']+'\n'+state['step5_d']+'\n'
    write_step6_llm = write_step6_template_d_EN | llm | StrOutputParser()
    for char in user_input:
        if '\u4e00' <= char <= '\u9fa5':
            write_step6_llm = write_step6_template_d | llm | StrOutputParser()
    step6_d = write_step6_llm.invoke(
        {'user_input': user_input,
         'finished_work': finished_work}
    )
    article = process_article(finished_work + step6_d)
    with open('review.md', 'w', encoding='utf-8') as file:
        # 将传入的字符串内容写入文件
        file.write(article)
    return {'step6_d': step6_d}



# 创建节点
workflow = StateGraph(Agent)
#workflow.add_node("judge_language", judge_language)
workflow.add_node("get_all_papers", get_all_papers)
workflow.add_node("get_survey_theme", get_survey_theme)
# workflow.add_node("judge_survey_type", judge_survey_type)

workflow.add_node("write_step1_a", write_step1_a)
workflow.add_node("write_step2_a", write_step2_a)
workflow.add_node("write_step3_a", write_step3_a)
workflow.add_node("write_step4_a", write_step4_a)
workflow.add_node("write_step5_a", write_step5_a)
workflow.add_node("write_step6_a", write_step6_a)

workflow.add_node("write_step1_b", write_step1_b)
workflow.add_node("write_step2_b", write_step2_b)
workflow.add_node("write_step3_b", write_step3_b)
workflow.add_node("write_step4_b", write_step4_b)
workflow.add_node("write_step5_b", write_step5_b)

workflow.add_node("write_step1_c", write_step1_c)
workflow.add_node("write_step2_c", write_step2_c)
workflow.add_node("write_step3_c", write_step3_c)
workflow.add_node("write_step4_c", write_step4_c)
workflow.add_node("write_step5_c", write_step5_c)

workflow.add_node("write_step1_d", write_step1_d)
workflow.add_node("write_step2_d", write_step2_d)
workflow.add_node("write_step3_d", write_step3_d)
workflow.add_node("write_step4_d", write_step4_d)
workflow.add_node("write_step5_d", write_step5_d)
workflow.add_node("write_step6_d", write_step6_d)

# 连接节点，形成工作流
workflow.set_entry_point("get_all_papers")
workflow.add_edge("get_all_papers", 'get_survey_theme')
# workflow.add_edge("get_survey_theme", 'judge_survey_type_1')
workflow.add_conditional_edges("get_survey_theme", judge_survey_type, {'a': 'write_step1_a', 'b': 'write_step1_b', 'c': 'write_step1_c', 'd': 'write_step1_d'})
workflow.add_edge("write_step1_a", 'write_step2_a')
workflow.add_edge("write_step1_a", 'write_step3_a')
workflow.add_edge("write_step1_a", 'write_step4_a')
workflow.add_edge("write_step1_a", 'write_step5_a')

workflow.add_edge("write_step2_a", 'write_step6_a')
workflow.add_edge("write_step3_a", 'write_step6_a')
workflow.add_edge("write_step4_a", 'write_step6_a')
workflow.add_edge("write_step5_a", 'write_step6_a')

workflow.add_edge("write_step1_b", 'write_step2_b')
workflow.add_edge("write_step1_b", 'write_step3_b')
workflow.add_edge("write_step1_b", 'write_step4_b')

workflow.add_edge("write_step2_b", 'write_step5_b')
workflow.add_edge("write_step3_b", 'write_step5_b')
workflow.add_edge("write_step4_b", 'write_step5_b')

workflow.add_edge("write_step1_c", 'write_step2_c')
workflow.add_edge("write_step1_c", 'write_step3_c')
workflow.add_edge("write_step1_c", 'write_step4_c')

workflow.add_edge("write_step2_c", 'write_step5_c')
workflow.add_edge("write_step3_c", 'write_step5_c')
workflow.add_edge("write_step4_c", 'write_step5_c')

workflow.add_edge("write_step1_d", 'write_step2_d')
workflow.add_edge("write_step1_d", 'write_step3_d')
workflow.add_edge("write_step1_d", 'write_step4_d')
workflow.add_edge("write_step1_d", 'write_step5_d')

workflow.add_edge("write_step2_d", 'write_step6_d')
workflow.add_edge("write_step3_d", 'write_step6_d')
workflow.add_edge("write_step4_d", 'write_step6_d')
workflow.add_edge("write_step5_d", 'write_step6_d')

app = workflow.compile()

if __name__ == "__main__":
    start_time = time.time()
    args = parse_args()
    config = {"recursion_limit": 500}
    inputs = {
        #只需要修改user_input的值来改变输入
        'user_input': args.topic
    }
    running_dict = {}
    for event in app.stream(inputs, config=config):
        for k, v in event.items():
            running_dict[k] = v
            if k != "__end__":
                print(v)
            print('----------' * 20)
    end_time = time.time()
    print(end_time-start_time)

